#!/usr/bin/env python
#-*- coding: UTF-8 -*-
# scheduler.py: scheduler
#
# Copyright (C) 2008-2010 Hily Jiang <hilyjiang@gmail.com>
#
# This file is part of dtask.
#
# This software is provided 'as-is', without any express or implied
# warranty.  In no event will the authors be held liable for any damages
# arising from the use of this software.
#
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not
#    claim that you wrote the original software. If you use this software
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
#    misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.

"""
Scheduler.
"""

import threading
import socket

from time import time, sleep

from task.scheduler import TaskContainer
from node.scheduler import Dispatcher


CFG_TASKSERVER_IP = '127.0.0.1'
CFG_TASKSERVER_PORT = 7709
CFG_NODESERVER_IP = '127.0.0.1'
CFG_NODESERVER_PORT = 8809
CFG_SCH_NAME = 'localsch'


class Scheduler(object):

    def __init__(self, taskServerIp, taskServerPort,
                 nodeServerIp, nodeServerPort,
                 schName=''):
        self.taskServerIp = taskServerIp
        self.taskServerPort = taskServerPort
        self.nodeServerIp = nodeServerIp
        self.nodeServerPort = nodeServerPort
        self.schName = schName
        self.taskContainer = TaskContainer()
        self.dispatcher = Dispatcher()
        self.connected = False

    def connect(self):
        """Connect to task server and node server."""
        self.connected = True
        if not self.taskContainer.connect(self.taskServerIp,
                            self.taskServerPort, self.schName):
            print 'connect to task server failed'
            self.connected = False
        if not self.dispatcher.connect(self.nodeServerIp,
                            self.nodeServerPort, self.schName):
            print 'connect to node server failed'
            self.connected = False
        return self.connected

    def run(self):
        """Run this scheduler."""
        if not self.connected:
            print 'please connect to server first'
            return False

        self.dispatchThread = threading.Thread(target=self.__dispatchThreadFunc)
        self.dispatchThread.setDaemon(True)
        self.dispatchThread.start()
        self.resultThread = threading.Thread(target=self.__resultThreadFunc)
        self.resultThread.setDaemon(True)
        self.resultThread.start()
        self.statusThread = threading.Thread(target=self.__statusThreadFunc)
        self.statusThread.setDaemon(True)
        self.statusThread.start()

        while True:
            sleep(10)

    def __dispatchThreadFunc(self):
        """Get tasks from task server and assign them to node server."""
        while True:
            # get task from task server
            taskInfo = self.taskContainer.getTask()
            if taskInfo == None: continue
            try:
                taskId, taskType, taskParams = taskInfo
                print 'task id:    ', taskId
                print 'task type:  ', taskType
                print 'task params:', taskParams
            except Exception, e:
                print e
                continue

            # dispatch task to node server
            self.dispatcher.dispatchTask(taskId, taskType, taskParams)

            # update status
            self.taskContainer.updateStatus(taskId, 'nodm %s %s %d 0' % \
                     (self.nodeServerIp, self.nodeServerPort, int(time())))

    def __resultThreadFunc(self):
        """Receive results from node server and send them to task server."""
        while True:
            # get result
            resultInfo = self.dispatcher.recvResult()
            try:
                taskId, result = resultInfo
            except Exception, e:
                print e
                continue

            print 'result of %s:' % resultInfo[0], resultInfo[1][:40]

            # update status
            self.taskContainer.updateStatus(taskId, 'nodm %s %s 0 %d' % \
                     (self.nodeServerIp, self.nodeServerPort, int(time())))

            # write result
            self.taskContainer.writeResult(taskId, result)

    def __statusThreadFunc(self):
        """Receive status from node server and send them to task server."""
        while True:
            # get status
            statusInfo = self.dispatcher.recvStatus()
            try:
                taskId, status = statusInfo
            except:
                continue

            # update status
            self.taskContainer.updateStatus(taskId, status)


def main():
    sch = Scheduler(CFG_TASKSERVER_IP, CFG_TASKSERVER_PORT,
                    CFG_NODESERVER_IP, CFG_NODESERVER_PORT,
                    CFG_SCH_NAME)
    if not sch.connect(): return
    sch.run()

if __name__ == '__main__':
    main()
